package com.github.plexpt.ws;


import cn.hutool.core.io.FileUtil;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.TypeReference;
import com.github.plexpt.chatgpt.Chatbot;
import com.github.plexpt.chatgpt.Start;
import com.github.plexpt.entity.Message;
import com.github.plexpt.entity.Params;
import com.github.plexpt.entity.Request;
import com.google.gson.Gson;
import lombok.SneakyThrows;

import javax.websocket.*;
import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;


/**
 * ws正向代理
 * gocq作为服务端
 * bot作为客户端
 */

@ClientEndpoint
public class Client {

    public static Session session;
    public static Client INSTANCE;
    public static final Gson gson = new Gson();

    @SneakyThrows
    public static String getInput(String prompt) {
        System.out.print(prompt);
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
        List<String> lines = new ArrayList<>();
        String line;
        try {
            while ((line = reader.readLine()) != null && !line.isEmpty()) {
                lines.add(line);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return lines.stream().collect(Collectors.joining("\n"));
    }


    public Client(String url) throws DeploymentException, IOException {
        session = ContainerProvider.getWebSocketContainer().connectToServer(this, URI.create(url));
    }

    public synchronized static boolean connect(String url) {
        try {
            INSTANCE = new Client(url);
            return true;
        } catch (DeploymentException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @OnOpen
    public void onOpen(Session session) {
        System.out.println("连接成功");
    }

    @OnMessage
    public void onMessage(String json) {
        //System.out.println(json);
        Message message = JSONObject.parseObject(json, Message.class);
        if ("message".equals(message.getPost_type())) {
            Request<Params> paramsRequest = new Request<>();
            paramsRequest.setAction("send_msg");
            Params params = new Params();
            params.setUser_id(message.getUser_id());
            params.setMessage(Chatbot(message.getMessage()));
            params.setMessage_type("private");
            params.setAuto_escape(true);
            paramsRequest.setParams(params);
            sendMessage(JSONObject.toJSONString(paramsRequest));
        }

    }

    public String Chatbot(String prompt) {
        StringBuilder stringBuilder = new StringBuilder();
        if (FileUtil.exist("config.json")) {

            String configString = FileUtil.readUtf8String(new File("config.json"));

            Map<String, String> params = JSON.parseObject(configString,
                    new TypeReference<Map<String, String>>() {
                    });

            Chatbot chatbot = new Chatbot(params, null);


            try {
                System.out.println("Chatbot: ");
                List<String> formattedParts = new ArrayList<>();
                Map<String, Object> message = chatbot.getChatResponse(prompt, "stream");
                // Split the message by newlines
                String[] messageParts = message.get("message").toString().split("\n");

                for (String part : messageParts) {
                    String[] wrappedParts = part.split("\n");
                    for (String wrappedPart : wrappedParts) {
                        stringBuilder.append(wrappedPart);
                    }
                }

            } catch (Exception e) {
                System.out.println("Something went wrong!");
                e.printStackTrace();
            }


        } else {
            System.out.println("Please create and populate config.json to continue");
            if (!FileUtil.exist("config.json")) {
                FileUtil.writeUtf8String("", "config.json");
            }
        }


        return stringBuilder.toString();

    }

    @OnClose
    public void onClose(Session session) {
        System.out.println("连接关闭");
    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        System.out.println("连接异常");
    }

    public static void sendMessage(String json) {
        Client.INSTANCE.session.getAsyncRemote().sendText(json);
    }


    @SneakyThrows
    public static void main(String[] args) {
        Start start = new Start();
        start.run();
        if (FileUtil.exist("config.json")) {

            String configString = FileUtil.readUtf8String(new File("config.json"));

            Map<String, String> params = JSON.parseObject(configString,
                    new TypeReference<Map<String, String>>() {
                    });

            Chatbot chatbot = new Chatbot(params, null);

            String prompt;
            while (true) {
                prompt = getInput("\nYou:\n");
                if (prompt.startsWith("!")) {
                    if (prompt.equals("!help")) {
                        System.out.println("\n!help - Show this message");
                        System.out.println("!reset - Forget the current conversation");
                        System.out.println("!refresh - Refresh the session authentication");
                        System.out.println("!rollback - Rollback the conversation by 1 message");
                        System.out.println("!config - Show the current configuration");
                        System.out.println("!exit - Exit the program");
                        continue;
                    } else if (prompt.equals("!reset")) {
                        chatbot.resetChat();
                        System.out.println("Chat session reset.");
                        continue;
                    } else if (prompt.equals("!refresh")) {
                        chatbot.refreshSession();
                        System.out.println("Session refreshed.\n");
                        continue;
                    } else if (prompt.equals("!rollback")) {
                        chatbot.rollbackConversation();
                        System.out.println("Chat session rolled back.");
                        continue;
                    } else if (prompt.equals("!config")) {
                        System.out.println(JSON.toJSONString(chatbot.getConfig()));
                        continue;
                    } else if (prompt.equals("!exit")) {
                        break;
                    }
                }

                if (Arrays.asList(args).contains("--text")) {
                    try {
                        System.out.println("Chatbot: ");
                        String out = "text";
                        Map<String, Object> message = chatbot.getChatResponse(prompt, out);
                        System.out.println(message.get("message"));
                    } catch (Exception e) {
                        System.out.println("Something went wrong!");
                        e.printStackTrace();
                        System.out.println(e);
                    }
                } else {

                    try {
                        System.out.println("Chatbot: ");
                        List<String> formattedParts = new ArrayList<>();
                        Map<String, Object> message = chatbot.getChatResponse(prompt, "stream");
                        // Split the message by newlines
                        String[] messageParts = message.get("message").toString().split("\n");

                        for (String part : messageParts) {
                            String[] wrappedParts = part.split("\n");
                            for (String wrappedPart : wrappedParts) {
                                System.out.println(wrappedPart);
                            }
                        }

                    } catch (Exception e) {
                        System.out.println("Something went wrong!");
                        e.printStackTrace();
                    }
                }

            }
        } else {
            System.out.println("Please create and populate config.json to continue");
            if (!FileUtil.exist("config.json")) {
                FileUtil.writeUtf8String("", "config.json");
            }
        }

    }


}

